import os
from ruamel.yaml import YAML
from typing import Tuple

class TaskGenerator:
    def __init__(self, base_yaml_path: str, work_dir: str):
        self.base_yaml_path = base_yaml_path
        self.work_dir = work_dir
        self.yaml = YAML()
        with open(base_yaml_path, 'r') as f:
            self.base_data = self.yaml.load(f)

    def _modify_yaml_path(self, data, path, value):
        # (此处省略之前讨论过的、用于修改嵌套字典的辅助函数)
        # ...
        keys = path.split('.')
        current_level = data
        for i, key in enumerate(keys):
            if i == len(keys) - 1:
                current_level[key] = value
            else:
                current_level = current_level[key]
        return data

    def generate(self, dimension_to_vary: str, value_a: str, value_b: str) -> Tuple[str, str]:
        """生成两个修改后的YAML文件，并返回它们的路径。"""
        # 生成配置A
        config_data_a = self.base_data.copy()
        self._modify_yaml_path(config_data_a, dimension_to_vary, value_a)
        
        # 生成配置B
        config_data_b = self.base_data.copy()
        self._modify_yaml_path(config_data_b, dimension_to_vary, value_b)
        
        # 保存文件
        base_name = os.path.splitext(os.path.basename(self.base_yaml_path))[0]
        path_a = os.path.join(self.work_dir, f"{base_name}_{value_a.replace('/', '_')}.yaml")
        path_b = os.path.join(self.work_dir, f"{base_name}_{value_b.replace('/', '_')}.yaml")
        
        with open(path_a, 'w') as f:
            self.yaml.dump(config_data_a, f)
            
        with open(path_b, 'w') as f:
            self.yaml.dump(config_data_b, f)
            
        return path_a, path_b
